#include <stdio.h>
#include <stdlib.h>
#include "PEStru.h"
#include "dos.h"

void CopyMemory(DWORD Dest, void *src, int bytes);

static int IsPE(FILE *fp)
{
	LONG PEMagic;
	IMAGE_DOS_HEADER DosHead;

	rewind(fp);
	fread(&DosHead, sizeof(IMAGE_DOS_HEADER), 1, fp);
	if (IMAGE_DOS_SIGNATURE != DosHead.e_magic)
		return 1;
	fseek(fp, DosHead.e_lfanew, SEEK_SET);
	fread(&PEMagic, sizeof(LONG), 1, fp);
	if (IMAGE_NT_SIGNATURE != PEMagic)
		return 2;
	return 0;
}

static void FixSectionHeader(WORD NumOfSections,
	PIMAGE_SECTION_HEADER pSectionHeader)
{
	register WORD i;
	for (i = 0; i < NumOfSections; i++)
	{
		if (pSectionHeader[i].Misc.VirtualSize
			> pSectionHeader[i].SizeOfRawData)
		{
			// The borland linker makes an mistake,
			// We just fix it :)
			DWORD tmp = pSectionHeader[i].SizeOfRawData;
			pSectionHeader[i].SizeOfRawData
				= pSectionHeader[i].Misc.VirtualSize;
			pSectionHeader[i].Misc.VirtualSize = tmp;
		}
	}
}

void LoadFSLimit4G(void);
void CopyMemory(DWORD Dest, void *src, int bytes);
static int LoadData(FILE *fp, DWORD pos, DWORD pAddr, DWORD Length)
{
	char buffer[512];
	fseek(fp, pos, 0);
	for (; Length > 512; Length -= 512)
	{
		fread(buffer, 512, 1, fp);
		CopyMemory(pAddr, (void *)buffer, 512);
		pAddr += 512L;
	}
	if (Length > 0)
	{
		fread(buffer, (size_t)Length, 1, fp);
		CopyMemory(pAddr, (void *)buffer, (int)Length);
	}
	return 0;
}

static int LoadSections(
	PIMAGE_FILE_HEADER		pImageHeader,
	PIMAGE_OPTIONAL_HEADER	pOptionalHeader,
	PIMAGE_SECTION_HEADER	pSectionHeader,
	FILE*							fp
	)
{
	register WORD i;
	DWORD Base = pOptionalHeader->ImageBase - 0xC0000000L;

	for (i = 0; i < pImageHeader->NumberOfSections; i++)
	{
		// Just for debug
		/*
		DWORD j = pSectionHeader[i].Misc.VirtualSize;
		printf("virtualsize:%lX, pointer to raw data: %lX, %lX\n", j,
		pSectionHeader[i].PointerToRawData,pSectionHeader[i].VirtualAddress + Base);
		*/

		if (LoadData(fp, pSectionHeader[i].PointerToRawData,
			pSectionHeader[i].VirtualAddress + Base,
			pSectionHeader[i].Misc.VirtualSize))
				return 1;
	}
	return 0;
}

WORD EnterProtectMode(
	DWORD							LinearAddress,
	PIMAGE_FILE_HEADER		pImageHeader,
	PIMAGE_OPTIONAL_HEADER	pOptionalHeader,
	PIMAGE_SECTION_HEADER	pSectionHeader
);

DWORD GetSysMemVol(void)
{
	DWORD TotalSysMem;

	asm 	mov	ax,0x0e801
	asm	int	0x15

	TotalSysMem = ((_EDX & 0xFFFF) << 6) + (DWORD)(_CX);
	return TotalSysMem;
}


void A20Enable(void);

extern char RegistryFileName[];

typedef struct
{
	DWORD TotalMemory;
	DWORD Flags;
	DWORD UnUsedHMABytes;
	DWORD UnUsedHMAAddr;
	DWORD RegistryPath;
}
DOSINFO, *PDOSINFO;

#define DOSFG_DOSINHMA			1

DWORD IsDosInHMA(void)
{
	// Check dos version
	unsigned char DosVer;
	unsigned char IsHMA;
	__asm {
		mov ah, 0x30
		int 0x21
		mov byte ptr DosVer,al
	}

	if (DosVer < 5) return 0L;
	__asm {
		mov ax, 3306
		int 0x21
		mov byte ptr IsHMA, dh
	}
	if (IsHMA & 0x10)
		return DOSFG_DOSINHMA;
	else
		return 0L;
}

void GetDosInfo(PDOSINFO pInfo)
{
	unsigned int bytes;
	unsigned int ves, vdi;

	pInfo->TotalMemory = GetSysMemVol();
	pInfo->Flags = IsDosInHMA();

		__asm {
			mov	ax, 0x4A01
			int	0x2F
			mov	bytes, bx
			mov	ax, es
			mov   ves, ax
			mov	vdi, di
		}

	 pInfo->UnUsedHMABytes = (DWORD)bytes;
	 pInfo->UnUsedHMAAddr =
		((((DWORD)ves) & 0xFFFF) * 16) + (((DWORD)(vdi)) & 0xFFFF);
	 if (bytes == 0)
		pInfo->Flags &= ~DOSFG_DOSINHMA;
	 else
		pInfo->Flags |= DOSFG_DOSINHMA;

	pInfo->RegistryPath = ((DWORD)(FP_SEG(RegistryFileName)) * 16);
	pInfo->RegistryPath += (DWORD)(FP_OFF(RegistryFileName));
   pInfo->RegistryPath += 0xC0000000;
}

int LoadPE(FILE *fp, char* szDLLName)
{
	WORD							RetVal;
	DOSINFO						DosInfo;
	IMAGE_FILE_HEADER			ImageHeader;
	IMAGE_OPTIONAL_HEADER	OptionalHeader;
	PIMAGE_SECTION_HEADER	pSectionHeader;

	if (IsPE(fp))
	{
		printf("Error in loading %s.\n", szDLLName);
		return 1;
	}
	fread(&ImageHeader, sizeof(IMAGE_FILE_HEADER), 1, fp);
	if (IMAGE_FILE_MACHINE_I386 != ImageHeader.Machine)
		return 2;
	if (!(IMAGE_FILE_EXECUTABLE_IMAGE
			& ImageHeader.Characteristics))
	{
		printf("%s is not a PE execute.\n", szDLLName);
		return 3;
	}
	if (ImageHeader.SizeOfOptionalHeader
		!= sizeof(IMAGE_OPTIONAL_HEADER))
	{
		printf("%s has been corrupted.\n", szDLLName);
		return 4;
	}
	fread(&OptionalHeader, sizeof(IMAGE_OPTIONAL_HEADER), 1, fp);
	if (IMAGE_NT_OPTIONAL_HDR_MAGIC != OptionalHeader.Magic)
	{
		printf("%s has a wrong optional header.\n", szDLLName);
		return 5;
	}
	if (0xC0100000L != OptionalHeader.ImageBase)
	{
		printf("%s is not a KERNEL DLL.\n", szDLLName);
		return 6;
	}
	pSectionHeader = (PIMAGE_SECTION_HEADER)
		malloc(sizeof(IMAGE_SECTION_HEADER) * ImageHeader.NumberOfSections);
	if (NULL == pSectionHeader)
	{
		printf("Not enough memory.\n");
		return 7;
	}
	fread(pSectionHeader, sizeof(IMAGE_SECTION_HEADER),
		ImageHeader.NumberOfSections, fp);
	FixSectionHeader(ImageHeader.NumberOfSections, pSectionHeader);
	if(LoadSections(&ImageHeader, &OptionalHeader, pSectionHeader, fp))
	{
		printf("Error in loading KERNEL DLL.\n");
		return 8;
	}

	A20Enable();

	GetDosInfo(&DosInfo);
	OptionalHeader.Reserved1 =
		((((DWORD)(_SS))& 0xFFFF) * 16L)
		+ (((DWORD)(&DosInfo)) & 0xFFFF);

	RetVal = EnterProtectMode(OptionalHeader.ImageBase
		+ OptionalHeader.AddressOfEntryPoint - 0xC0000000L,
		&ImageHeader, &OptionalHeader, pSectionHeader);

	if (RetVal == 0)
		printf("Return to DOS successfully.\n");
	else if (RetVal == 1)
		printf("Restarting system, please wait...\n");
	else
	{
		printf("Return to DOS because of an error.\n");
		printf("Error code is %d.\n", RetVal);
	}

	free(pSectionHeader);
	return 0;
}